from mindspore import context
from mindspore.common import set_seed
from mindspore.train.model import Model
from mindspore.train.serialization import load_checkpoint, load_param_into_net
import mindspore.dataset as ds
import argparse
from ms_losses import MultiLosses
from ms_model_abinet_iter import ABINetIterModel
from ms_utils import Config
import mindspore as ms
#from src.loss import CTCLoss
##from src.dataset import create_dataset
from ms_create_dataset_for_test import image_dataset, ms_ImageDataset
#from src.crnn import crnn
#from src.metric import CRNNAccuracy
from ms_metrics import TextAccuracy
#from src.model_utils.moxing_adapter import moxing_wrapper
#from src.model_utils.config import config
#from src.model_utils.device_adapter import get_device_id

import mindspore.nn as nn

import faulthandler
faulthandler.enable()

set_seed(1)


#context.set_context(mode=context.PYNATIVE_MODE,device_target='GPU',save_graphs=False)

context.set_context(mode=context.PYNATIVE_MODE, pynative_synchronize=True,device_target='GPU',save_graphs=False)


# @moxing_wrapper(pre_process=None)
def crnn_eval(config):
    #if config.device_target == 'Ascend':
        #device_id = get_device_id()
        #context.set_context(device_id=device_id)

    
    max_text_length = config.max_text_length
    # input_size = config.input_size
    # create dataset
    # dataset = create_dataset(name=config.eval_dataset,
    #                          dataset_path=config.eval_dataset_path,
    #                          batch_size=config.batch_size,
    #                          is_training=False,
    #                          config=config)
    dataset_train_image_generator = ms_ImageDataset(path='/home/data4/zyh/ABINet/data/evaluation/SVTP')
    #dataset_image_generator = ms_ImageDataset(path='/home/data4/zyh/ABINet/data/evaluation/CUTE80')


    image_dataset = ds.GeneratorDataset(dataset_train_image_generator, column_names= ['image','label','length'], shuffle=False)
    image_dataset = image_dataset.batch(batch_size= config.batch_size, drop_remainder=True)
    dataset = image_dataset
    # step_size = dataset.get_dataset_size()
    # loss = CTCLoss(max_sequence_length=config.num_step,
    #                max_label_length=max_text_length,
    #                batch_size=config.batch_size)
    loss = MultiLosses()

    #net = crnn(config, full_precision=config.device_target != 'Ascend')
    abinet = ABINetIterModel(config=config)
    # load checkpoint

    #param_dict = load_checkpoint(config.checkpoint_path)#暂时没有ckpt，故未添加
    #load_param_into_net(net, param_dict)#暂时没有ckpt，故未添加
    eval_net = CustomWithEvalCell(abinet)
    abinet.set_train(False)
    # define model
    # model = Model(network=abinet, eval_network=eval_net,
    #              metrics={'TextAccuracy': TextAccuracy(charset_path=config.dataset_charset_path,
    #          max_length=config.dataset_max_length + 1,
    #          case_sensitive=config.dataset_eval_case_sensisitves,
    #          model_eval=config.model_eval)})
    metric = TextAccuracy(charset_path=config.dataset_charset_path,
              max_length=config.dataset_max_length + 1,
              case_sensitive=config.dataset_eval_case_sensisitves,
              model_eval=config.model_eval)

    # 加载权重的参数
    ckpt_file_name = "/home/data4/zyh/ABINet/ABINet_Mindspore/abinet_ms/ms_abinet_vision.ckpt"
    param_dict = ms.load_checkpoint(ckpt_file_name)
    param_not_load = ms.load_param_into_net(abinet, param_dict)



    model = Model(abinet,loss)
    metric.clear()
    for batch_idx, (data, label,length) in enumerate(dataset):
        print(batch_idx)
        output = model.predict(data)
        metric.update(output, label)
    results = metric.eval()
    print(results)
    #result = model.eval(image_dataset,dataset_sink_mode=False)
    #print(result)

    # model = Model(abinet, 
    # loss_fn=loss, 
    # metrics={'TextAccuracy': TextAccuracy(charset_path=config.dataset_charset_path,
    #         max_length=config.dataset_max_length + 1,
    #         case_sensitive=config.dataset_eval_case_sensisitves,
    #         model_eval=config.model_eval)}
    #         )#还没加优化器!!!!!!
    # # start evaluation

    # res = model.eval(dataset, dataset_sink_mode=config.device_target == 'GPU')
    # print("result:", res, flush=True)



class CustomWithEvalCell(nn.Cell):
    """自定义多标签评估网络"""

    def __init__(self, network):
        super(CustomWithEvalCell, self).__init__(auto_prefix=False)
        self.network = network

    def construct(self, data, label1, length):
        output = self.network(data)
        return output, label1, length












def main():
    set_seed(1)
    context.set_context(mode=context.PYNATIVE_MODE, save_graphs=False)
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='/home/data/zhangyh22/ABINet/ABINet_Mindspore/configs/train_abinet.yaml',
                        help='path to config file')
    parser.add_argument('--input', type=str, default='figs/test')
    parser.add_argument('--cuda', type=int, default=-1)
    parser.add_argument('--checkpoint', type=str, default='workdir/train-abinet/best-train-abinet.pth')
    parser.add_argument('--model_eval', type=str, default='alignment', 
                        choices=['alignment', 'vision', 'language'])
    parser.add_argument('--batch_size', type=int, default=1)
    parser.add_argument('--max_text_length', type=int, default=30)
    args = parser.parse_args()
    config = Config(args.config)
    if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
    if args.model_eval is not None: config.model_eval = args.model_eval
    config.global_phase = 'test'
    config.model_vision_checkpoint, config.model_language_checkpoint = None, None
    device = 'cpu' if args.cuda < 0 else f'cuda:{args.cuda}'
    config.batch_size = 72


    crnn_eval(config)







if __name__ == '__main__':
    main()